#!/usr/bin/env python3
"""
数据库迁移脚本 - 更新客户表结构
添加新字段：客户分类、归属地、来源等，并移除旧字段
"""

import sqlite3
import os
from datetime import datetime

# 数据库文件路径
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'database', 'eqsystem.db')

def migrate_customer_table():
    """迁移客户表结构"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    try:
        print("开始迁移客户表结构...")
        
        # 1. 检查表是否存在
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='customer'")
        if not cursor.fetchone():
            print("错误：客户表不存在")
            return False
        
        # 2. 创建临时表，包含新结构
        print("创建临时表...")
        cursor.execute('''
            CREATE TABLE customer_temp (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                name VARCHAR(200) NOT NULL,
                customer_number VARCHAR(50) UNIQUE NOT NULL,
                category VARCHAR(50) NOT NULL DEFAULT '普通客户',
                category_other VARCHAR(100),
                location VARCHAR(100) NOT NULL,
                source VARCHAR(50) NOT NULL DEFAULT '网站',
                source_other VARCHAR(100),
                address TEXT,
                province_id INTEGER,
                city_id INTEGER,
                district_id INTEGER,
                manager_id INTEGER,
                converted_from_lead_id INTEGER,
                conversion_date DATETIME,
                conversion_reason TEXT,
                FOREIGN KEY (manager_id) REFERENCES user (id),
                FOREIGN KEY (converted_from_lead_id) REFERENCES lead (id),
                FOREIGN KEY (province_id) REFERENCES regions (id),
                FOREIGN KEY (city_id) REFERENCES regions (id),
                FOREIGN KEY (district_id) REFERENCES regions (id)
            )
        ''')
        
        # 3. 迁移现有数据
        print("迁移现有数据...")
        cursor.execute('''
            INSERT INTO customer_temp (
                id, created_at, updated_at, name, customer_number, 
                category, location, source, address, manager_id
            )
            SELECT 
                id, created_at, updated_at, name, customer_number,
                '普通客户', '本地', '网站', address, manager_id
            FROM customer
        ''')
        
        # 4. 删除旧表
        print("删除旧表...")
        cursor.execute("DROP TABLE customer")
        
        # 5. 重命名临时表
        print("重命名表...")
        cursor.execute("ALTER TABLE customer_temp RENAME TO customer")
        
        # 6. 重建索引
        print("重建索引...")
        cursor.execute("CREATE INDEX ix_customer_manager_id ON customer (manager_id)")
        cursor.execute("CREATE INDEX ix_customer_customer_number ON customer (customer_number)")
        
        # 提交事务
        conn.commit()
        print("客户表结构迁移完成！")
        
        return True
        
    except Exception as e:
        print(f"迁移过程中发生错误: {str(e)}")
        conn.rollback()
        return False
        
    finally:
        conn.close()

def check_regions_table():
    """检查regions表是否存在，如果不存在则创建基础数据"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    try:
        # 检查表是否存在
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='regions'")
        if not cursor.fetchone():
            print("创建regions表...")
            cursor.execute('''
                CREATE TABLE regions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    name VARCHAR(50) NOT NULL,
                    code VARCHAR(20) UNIQUE NOT NULL,
                    level VARCHAR(10) NOT NULL,
                    parent_id INTEGER,
                    FOREIGN KEY (parent_id) REFERENCES regions (id)
                )
            ''')
            
            # 插入基础省份数据
            provinces = [
                ('110000', '北京', '省'),
                ('120000', '天津', '省'),
                ('130000', '河北省', '省'),
                ('140000', '山西省', '省'),
                ('150000', '内蒙古自治区', '省'),
                ('210000', '辽宁省', '省'),
                ('220000', '吉林省', '省'),
                ('230000', '黑龙江省', '省'),
                ('310000', '上海', '省'),
                ('320000', '江苏省', '省'),
                ('330000', '浙江省', '省'),
                ('340000', '安徽省', '省'),
                ('350000', '福建省', '省'),
                ('360000', '江西省', '省'),
                ('370000', '山东省', '省'),
                ('410000', '河南省', '省'),
                ('420000', '湖北省', '省'),
                ('430000', '湖南省', '省'),
                ('440000', '广东省', '省'),
                ('450000', '广西壮族自治区', '省'),
                ('460000', '海南省', '省'),
                ('500000', '重庆', '省'),
                ('510000', '四川省', '省'),
                ('520000', '贵州省', '省'),
                ('530000', '云南省', '省'),
                ('540000', '西藏自治区', '省'),
                ('610000', '陕西省', '省'),
                ('620000', '甘肃省', '省'),
                ('630000', '青海省', '省'),
                ('640000', '宁夏回族自治区', '省'),
                ('650000', '新疆维吾尔自治区', '省')
            ]
            
            cursor.executemany(
                "INSERT INTO regions (code, name, level) VALUES (?, ?, ?)",
                provinces
            )
            
            conn.commit()
            print("regions表创建完成，已插入基础省份数据")
        
        return True
        
    except Exception as e:
        print(f"检查regions表时发生错误: {str(e)}")
        return False
        
    finally:
        conn.close()

if __name__ == "__main__":
    print(f"开始数据库迁移，时间: {datetime.now()}")
    print(f"数据库路径: {DB_PATH}")
    
    # 检查regions表
    if not check_regions_table():
        print("检查regions表失败，终止迁移")
        exit(1)
    
    # 迁移客户表
    if migrate_customer_table():
        print("数据库迁移成功完成！")
    else:
        print("数据库迁移失败！")
        exit(1)